package gorm_pool

import (
	"context"
	"errors"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/jinzhu/gorm"

	"gorm_pool/databases"
	"gorm_pool/databases/mysql"
)

// 定义了一系列错误，用于表示连接池操作过程中可能出现的不同错误情况
var (
	// 当尝试从已关闭的连接池中获取连接时返回该错误
	ErrGetFromClosedPool = errors.New("pool: get from closed pool")
	// 当从连接池获取连接时超过了最大等待时间返回该错误
	ErrExceedingMaxWaitingDuration = errors.New("pool: exceeding the maximum waiting duration")
	// 当尝试连接不被连接池支持的 SQL 类型时返回该错误
	ErrSQLType = errors.New("pool: sql type does not support")
	// 当 KeepConn（保持的空闲连接数）大于连接池容量时返回该错误
	ErrKeepLTCapacity = errors.New("pool: KeepConn larger than Capacity")
	// 当提供的连接池容量无效（为 0）时返回该错误
	ErrCapacity = errors.New("pool: invalid capacity size")
	// 当数据库连接参数为空时返回该错误
	ErrEmptyArgs = errors.New("pool: args cannot be empty")
)

// Options 结构体用于配置连接池的参数
type Options struct {
	// 数据库类型，例如 MySQL、SQLite3 等
	Type databases.SQLType
	// 数据库连接配置参数
	Args interface{}
	// 生成连接键的函数
	KeyFunc func() string
	// 当没有工作要做时，要保持的空闲连接数，该值应小于 Capacity
	KeepConn uint64
	// 在给定时间内连接池分配的最大连接数
	Capacity uint64
	// 从连接池获取连接的最大等待时间
	MaxWaitDuration time.Duration
	// 数据库连接的创建器，用于创建实际的数据库连接
	connector databases.Connector
}

// validate 方法用于验证连接池配置选项的有效性
func (o *Options) validate() (err error) {
	// 根据数据库类型选择对应的连接创建器，目前仅支持 MySQL
	switch o.Type {
	case databases.MySQL:
		o.connector = mysql.New(o.Args)
	default:
		return ErrSQLType
	}

	// 检查数据库连接参数是否为空
	if o.Args == nil {
		return ErrEmptyArgs
	}

	// 检查连接池容量是否为 0
	if o.Capacity == 0 {
		return ErrCapacity
	}

	// 检查保持的空闲连接数是否大于连接池容量
	if o.KeepConn > o.Capacity {
		return ErrKeepLTCapacity
	}

	return nil
}

// Pool 结构体表示数据库连接池
type Pool struct {
	// 数据库类型
	Type databases.SQLType
	// 数据库连接配置参数
	Args interface{}
	// 生成连接键的函数
	keyFunc func() string
	// 数据库连接的创建器
	connector databases.Connector
	// 当没有工作要做时，要保持的空闲连接数
	keepConn uint64
	// 连接池的最大连接数
	capacity uint64
	// 从连接池获取连接的最大等待时间
	maxWaitDuration time.Duration
	// 互斥锁，用于保护连接池的状态和相关操作
	mu sync.Mutex
	// 存储空闲连接的列表
	idleConn *conns
	// 存储繁忙连接的列表
	busyConn *conns
	// 连接池是否已关闭的标志
	closed bool
	// 用于限制打开连接数量的通道，当 p.Wait 为 true 时起作用
	ch chan struct{}
	// 当前正在等待获取连接的连接数
	currentWaitCount uint64
	// 总共等待获取连接的连接数
	totalWaitCount uint64
	// 等待新连接的总时间
	waitDuration time.Duration
	// 由于超过最大等待时间而放弃获取连接的次数计数器
	droppedGetCount uint64
	// 上下文，用于控制连接池的生命周期
	ctx context.Context
}

// NewPool 函数用于创建一个新的连接池实例
func NewPool(ctx context.Context, opts Options) (p *Pool, err error) {
	// 验证连接池配置选项的有效性
	err = opts.validate()
	if err != nil {
		return nil, err
	}

	// 初始化连接池的各个属性
	p = &Pool{
		Type:             opts.Type,
		Args:             opts.Args,
		keyFunc:          opts.KeyFunc,
		connector:        opts.connector,
		keepConn:         opts.KeepConn,
		capacity:         opts.Capacity,
		maxWaitDuration:  opts.MaxWaitDuration,
		mu:               sync.Mutex{},
		idleConn:         newConns(),
		busyConn:         newConns(),
		closed:           false,
		ch:               make(chan struct{}, opts.Capacity),
		currentWaitCount: 0,
		totalWaitCount:   0,
		waitDuration:     0,
		droppedGetCount:  0,
		ctx:              ctx,
	}

	// 如果要保持的空闲连接数为 0，检查连接参数是否能创建有效连接
	if p.keepConn == 0 {
		err = p.checkArgs()
		if err != nil {
			return nil, err
		}
		return p, nil
	}

	// 初始化指定数量的空闲连接
	for i := uint64(0); i < p.keepConn; i++ {
		err = p.initConn()
		if err != nil {
			return nil, err
		}
	}

	// 启动一个 goroutine 监听上下文的取消信号，当信号到来时关闭连接池
	go func() {
		select {
		case <-p.ctx.Done():
			p.Close()
		default:
		}
	}()
	return p, nil
}

// initConn 方法用于初始化一个新的数据库连接并添加到空闲连接列表中
func (p *Pool) initConn() error {
	var (
		key string
		db  *gorm.DB
		err error
	)

	// 使用连接创建器打开一个数据库连接
	db, err = p.connector.Open()
	if err != nil {
		return err
	}

	// 如果没有提供生成连接键的函数，则使用 UUID 生成一个唯一键
	if p.keyFunc == nil {
		key = strings.ReplaceAll(uuid.New().String(), "-", "")
	} else {
		key = p.keyFunc()
	}

	// 创建一个连接对象并设置相关属性
	conn := &Conn{
		DB:           db,
		Key:          key,
		Created:      time.Now(),
		Updated:      time.Now(),
		UsageCounter: 0,
	}
	// 将连接添加到限制打开连接数量的通道中，并将连接放入空闲连接列表
	p.ch <- struct{}{}
	p.idleConn.put(conn)
	return nil
}

// checkArgs 方法用于检查连接参数是否能创建一个有效的数据库连接
func (p *Pool) checkArgs() error {
	// 使用连接创建器打开一个数据库连接
	db, err := p.connector.Open()
	if err != nil {
		return err
	}
	// 关闭连接，仅用于验证连接参数的有效性
	db.Close()
	return nil
}

// Get 方法用于从连接池获取一个数据库连接
func (p *Pool) Get() (conn *Conn, err error) {
	// 检查上下文是否已取消，如果是则增加放弃获取连接的次数并返回错误
	select {
	case <-p.ctx.Done():
		p.mu.Lock()
		p.droppedGetCount++
		p.mu.Unlock()
		return nil, ErrGetFromClosedPool
	default:
		// 加锁以保护连接池状态
		p.mu.Lock()
		// 检查连接池是否已关闭，如果是则增加放弃获取连接的次数并返回错误
		if p.closed {
			p.droppedGetCount++
			p.mu.Unlock()
			return nil, ErrGetFromClosedPool
		}

		// 如果有空闲连接，从空闲连接列表中获取一个连接并从通道中移除一个信号
		if p.idleConn.size > 0 {
			conn = p.get()
			<-p.ch
			p.mu.Unlock()
			return conn, nil
		}

		// 如果繁忙连接数小于连接池容量，创建一个新连接并获取
		if p.busyConn.size < p.capacity {
			err = p.initConn()
			if err != nil {
				return nil, err
			}
			conn = p.get()
			<-p.ch
			p.mu.Unlock()
			return conn, nil
		}

		// 如果繁忙连接数达到容量，设置一个定时器等待连接可用
		timer := time.NewTimer(p.maxWaitDuration)
		start := time.Now()
		p.currentWaitCount++
		p.totalWaitCount++
		p.mu.Unlock()
		// 等待上下文取消、连接可用或定时器超时
		select {
		case <-p.ctx.Done():
			p.mu.Unlock()
			return nil, ErrGetFromClosedPool
		case <-p.ch:
			p.mu.Lock()
			conn = p.get()
			p.waitDuration += time.Since(start)
			p.currentWaitCount--
			p.mu.Unlock()
			return conn, nil
		case <-timer.C:
			p.mu.Lock()
			p.waitDuration += time.Since(start)
			p.droppedGetCount++
			p.currentWaitCount--
			p.mu.Unlock()
			return nil, ErrExceedingMaxWaitingDuration
		}
	}
}

// get 方法用于从空闲连接列表中获取一个连接并将其移动到繁忙连接列表
func (p *Pool) get() (conn *Conn) {
	conn = p.idleConn.get()
	p.busyConn.put(conn)
	return conn
}

// Put 方法用于将一个连接放回连接池
func (p *Pool) Put(conn *Conn) {
	// 加锁以保护连接池状态
	p.mu.Lock()
	// 从繁忙连接列表中移除该连接
	p.busyConn.deleteByKey(conn.Key)

	// 如果空闲连接数小于要保持的空闲连接数且连接池未关闭，将连接放回空闲连接列表并向通道中添加一个信号
	if p.idleConn.size < p.keepConn && !p.closed {
		p.idleConn.put(conn)
		p.ch <- struct{}{}
	} else {
		// 否则关闭连接并将其设置为 nil
		conn.DB.Close()
		conn = nil
	}
	// 解锁
	p.mu.Unlock()
}

// Close 方法用于关闭连接池，将连接池标记为已关闭并关闭所有空闲连接
func (p *Pool) Close() {
	p.mu.Lock()
	p.closed = true
	p.idleConn.close()
	p.mu.Unlock()
}

// Status 方法用于获取连接池的当前状态信息
func (p *Pool) Status() (ps PoolState) {
	// 加锁以保护连接池状态
	p.mu.Lock()
	// 填充连接池状态信息结构体
	ps = PoolState{
		IdleConnsState: ConnsState{
			Size:  p.idleConn.size,
			Conns: p.idleConn.conns,
		},
		BusyConnsState: ConnsState{
			Size:  p.busyConn.size,
			Conns: p.busyConn.conns,
		},
		Capacity:             p.capacity,
		Closed:               p.closed,
		Size:                 p.busyConn.size + p.idleConn.size,
		TotalWaitingDuration: p.waitDuration,
		CurrentWaitCount:     p.currentWaitCount,
		TotalWaitCount:       p.totalWaitCount,
		DroppedGetCount:      p.droppedGetCount,
	}
	// 解锁
	p.mu.Unlock()
	return ps
}

// conns 结构体用于管理连接的列表
type conns struct {
	// 连接键的列表
	keys []string
	// 连接对象的映射，键为连接键，值为连接对象
	conns map[string]*Conn
	// 连接的数量
	size uint64
}

// Conn 结构体用于封装内部的数据库连接，并记录连接的相关信息
type Conn struct {
	// 实际的数据库连接对象
	DB *gorm.DB
	// 连接的唯一键
	Key string
	// 连接的创建时间
	Created, Updated time.Time
	// 连接的使用次数
	UsageCounter uint64
}

// newConns 函数用于创建一个新的连接列表实例
func newConns() *conns {
	return &conns{
		keys:  []string{},
		conns: map[string]*Conn{},
		size:  0,
	}
}

// get 方法用于从连接列表中获取一个连接，并更新相关信息
func (cs *conns) get() (conn *Conn) {
	// 获取连接键并从列表中移除
	key := cs.keys[0]
	cs.keys = cs.keys[1:]
	// 获取连接对象并从映射中删除
	conn = cs.conns[key]
	delete(cs.conns, key)
	cs.size--
	// 增加连接的使用次数并更新连接的更新时间
	conn.UsageCounter++
	conn.Updated = time.Now()
	return conn
}

// put 方法用于将一个连接添加到连接列表中，并更新相关信息
func (cs *conns) put(conn *Conn) {
	// 将连接键添加到列表中
	cs.keys = append(cs.keys, conn.Key)
	// 将连接对象添加到映射中
	cs.conns[conn.Key] = conn
	cs.size++
	// 更新连接的更新时间
	conn.Updated = time.Now()
}

// deleteByKey 方法用于根据连接键从连接列表中删除一个连接，并更新相关信息
func (cs *conns) deleteByKey(key string) {
	// 创建一个新的连接键列表
	keys := make([]string, cs.size-1)
	for _, val := range cs.keys {
		if val != key {
			keys = append(keys, val)
		} else {
			// 更新连接的更新时间
			cs.conns[val].Updated = time.Now()
		}
	}
	cs.keys = keys
	cs.size--
	// 从映射中删除连接对象
	delete(cs.conns, key)
}

// close 方法用于关闭连接列表中的所有连接，并清空连接列表和映射
func (cs *conns) close() {
	for _, conn := range cs.conns {
		conn.DB.Close()
	}
	for _, key := range cs.keys {
		delete(cs.conns, key)
	}
}

// ConnsState 结构体用于表示连接列表的状态，包括连接数量和连接对象的映射
type ConnsState struct {
	Size  uint64
	Conns map[string]*Conn
}

// PoolState 结构体用于表示连接池的状态，包含空闲连接状态、繁忙连接状态、容量、是否关闭等信息
type PoolState struct {
	IdleConnsState       ConnsState
	BusyConnsState       ConnsState
	Capacity             uint64
	Size                 uint64
	Closed               bool
	TotalWaitingDuration time.Duration
	CurrentWaitCount     uint64
	TotalWaitCount       uint64
	DroppedGetCount      uint64
}
